import numpy as np
from .emittance_hanlder import calculate_rms_emittance

def calculate_twiss_parameters(x, x_prime):
    """
    Calculate Twiss parameters (β, α, γ) from phase space coordinates.
    Uses the correct RMS emittance calculation from emittance_hanlder.
    """
    # Use the standardized RMS emittance calculation
    emittance = calculate_rms_emittance(x, x_prime)
    
    if emittance > 1e-15:
        # Remove means to center the phase space ellipse
        x_centered = x - np.mean(x)
        x_prime_centered = x_prime - np.mean(x_prime)
        
        # Calculate second order moments
        sigma_11 = np.mean(x_centered**2)      # ⟨x²⟩
        sigma_12 = np.mean(x_centered * x_prime_centered)  # ⟨xx'⟩
        sigma_22 = np.mean(x_prime_centered**2)  # ⟨x'²⟩
        
        # Calculate Twiss parameters
        beta = sigma_11 / emittance      # β = ⟨x²⟩/ε
        alpha = -sigma_12 / emittance    # α = -⟨xx'⟩/ε
        gamma = sigma_22 / emittance     # γ = ⟨x'²⟩/ε
        
        # Verify Courant-Snyder invariant: γ = (1 + α²)/β
        gamma_expected = (1 + alpha**2) / beta
        consistency = np.abs(gamma - gamma_expected) / gamma_expected
        
        if consistency > 0.01:
            print(f"Warning: Twiss parameter consistency check failed: {consistency:.4f}")
    else:
        beta, alpha, gamma = 1.0, 0.0, 1.0
    
    return beta, alpha, gamma, emittance
